#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：MicInter_IPCamera 
@File    ：测试.py
@IDE     ：PyCharm 
@Author  ：董文捷
@Date    ：2025/11/5 13:57 
'''
import cv2
import numpy as np
import os
from ultralytics import YOLO
import time
from tqdm import tqdm
from scipy.signal import find_peaks
import matplotlib.pyplot as plt


class ThroughputDetector:
    """吞吐检测器 - 基于中线交点法"""

    def __init__(self, model_path='runs/detect/train/weights/best.pt'):
        """
        初始化吞吐检测器（中线交点法）
        """
        print(f"🔍 加载训练好的模型: {model_path}")

        if not os.path.exists(model_path):
            raise FileNotFoundError(f"❌ 找不到模型文件: {model_path}")

        try:
            self.model = YOLO(model_path)
            print("✅ 模型加载成功")
        except Exception as e:
            raise RuntimeError(f"❌ 模型加载失败: {str(e)}")

        # 配置参数
        self.confidence_threshold = 0.1

        # 吞吐检测状态
        self.brightness_values = []  # 存储所有亮度值
        self.frame_numbers = []  # 存储对应的帧号
        self.crossing_points = []  # 存储中线交点位置
        self.current_throughput = 0  # 当前吞吐数量
        self.crossing_events = []  # 存储交点事件

        print("✅ 中线交点法吞吐器初始化完成")

    def get_center_region_brightness(self, frame, center_point, radius=20):
        """
        获取中心区域的亮度值
        """
        x, y = int(center_point[0]), int(center_point[1])
        h, w = frame.shape[:2]

        # 确保区域在图像范围内
        x1 = max(0, x - radius)
        y1 = max(0, y - radius)
        x2 = min(w, x + radius)
        y2 = min(h, y + radius)

        if x1 >= x2 or y1 >= y2:
            return 0

        # 提取区域并计算亮度
        region = frame[y1:y2, x1:x2]
        if len(region) == 0:
            return 0

        if len(region.shape) == 3:
            region = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)

        brightness = np.mean(region) / 255.0
        return brightness

    def detect_center_point(self, frame):
        """
        检测帧中的中心点
        """
        try:
            results = self.model(frame, imgsz=640, conf=self.confidence_threshold, verbose=False)

            for result in results:
                if hasattr(result, 'boxes') and result.boxes is not None and len(result.boxes) > 0:
                    boxes = result.boxes.xyxy.cpu().numpy()
                    confidences = result.boxes.conf.cpu().numpy()

                    # 取置信度最高的检测框
                    if len(confidences) > 0:
                        max_conf_idx = np.argmax(confidences)
                        box = boxes[max_conf_idx]
                        conf = confidences[max_conf_idx]

                        # 计算中心点
                        x_center = (box[0] + box[2]) / 2
                        y_center = (box[1] + box[3]) / 2

                        return (x_center, y_center), conf

            return None, 0

        except Exception as e:
            print(f"⚠️ 中心点检测错误: {str(e)}")
            return None, 0

    def detect_midline_crossings(self):
        """
        检测亮度曲线与中线的交点
        返回：交点数量（吞吐次数 = 有效交点对数）
        过滤单个交点：只计算成对的交点（一个上升沿和一个下降沿）
        """
        if len(self.brightness_values) < 2:
            return 0, []

        # 计算中线（亮度平均值）
        midline = np.mean(self.brightness_values)

        # 检测所有可能的交点
        all_crossings = []
        for i in range(1, len(self.brightness_values)):
            prev_val = self.brightness_values[i - 1]
            curr_val = self.brightness_values[i]

            # 检查是否穿过中线
            if (prev_val < midline and curr_val >= midline) or \
                    (prev_val > midline and curr_val <= midline):
                all_crossings.append(i)

        # 过滤单个交点 - 只保留成对的交点
        valid_crossings = []
        i = 0
        while i < len(all_crossings) - 1:
            # 检查当前交点和下一个交点是否构成一对（一个上升一个下降）
            current_idx = all_crossings[i]
            next_idx = all_crossings[i + 1]

            # 获取交点的亮度值
            current_val = self.brightness_values[current_idx]
            next_val = self.brightness_values[next_idx]

            # 检查是否构成一对（一个高于中线，一个低于中线）
            if (current_val >= midline and next_val <= midline) or \
                    (current_val <= midline and next_val >= midline):
                valid_crossings.append(current_idx)
                valid_crossings.append(next_idx)
                i += 2  # 跳过下一对
            else:
                i += 1  # 移动到下一个点

        # 吞吐次数 = 有效交点对数
        throughput_count = len(valid_crossings) // 2

        return throughput_count, valid_crossings, midline

    def update_throughput_count(self, frame_idx, brightness, midline):
        """
        实时更新吞吐数量
        """
        if len(self.brightness_values) < 2:
            return self.current_throughput

        # 检查是否穿过中线
        prev_val = self.brightness_values[-2] if len(self.brightness_values) >= 2 else brightness
        curr_val = brightness

        # 检查是否穿过中线
        crossed = False
        direction = None

        if prev_val < midline and curr_val >= midline:
            crossed = True
            direction = "up"
        elif prev_val > midline and curr_val <= midline:
            crossed = True
            direction = "down"

        if crossed:
            # 记录交点事件
            self.crossing_events.append({
                'frame': frame_idx,
                'brightness': curr_val,
                'direction': direction,
                'time': time.time()
            })

            # 检查是否形成有效对
            if len(self.crossing_events) >= 2:
                last_event = self.crossing_events[-1]
                prev_event = self.crossing_events[-2]

                # 如果方向相反且时间间隔合理（避免误检）
                if (last_event['direction'] != prev_event['direction'] and
                        last_event['frame'] - prev_event['frame'] < 100):  # 最大100帧间隔
                    self.current_throughput += 1
                    # 移除已配对的事件
                    self.crossing_events = self.crossing_events[:-2]

        return self.current_throughput

    def draw_detection_info(self, frame, center_point, brightness, throughput_count):
        """
        在帧上绘制检测信息和吞吐数量
        """
        display_frame = frame.copy()

        # 绘制中心点
        if center_point:
            x, y = int(center_point[0]), int(center_point[1])
            cv2.circle(display_frame, (x, y), 10, (0, 255, 0), 2)
            cv2.circle(display_frame, (x, y), 20, (0, 255, 0), 1)

        # 绘制吞吐数量
        cv2.putText(display_frame, f"Throughput: {throughput_count}", (20, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # 绘制亮度值
        cv2.putText(display_frame, f"Brightness: {brightness:.3f}", (20, 90),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

        # 绘制帧信息
        cv2.putText(display_frame, f"Frames: {len(self.brightness_values)}", (20, 130),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)

        return display_frame

    def plot_analysis(self, video_name):
        """
        绘制中线交点分析图表
        """
        if len(self.brightness_values) < 2:
            return None, 0, 0

        # 检测中线交点
        throughput_count, crossings, midline = self.detect_midline_crossings()

        # 绘制图表
        plt.figure(figsize=(12, 8))

        frames = range(len(self.brightness_values))

        # 绘制亮度曲线
        plt.plot(frames, self.brightness_values, 'b-', label='Brightness', alpha=0.8, linewidth=2)

        # 绘制中线
        plt.axhline(y=midline, color='r', linestyle='-', alpha=0.7, linewidth=2,
                    label=f'Midline: {midline:.3f}')

        # 绘制交点
        if crossings:
            crossing_values = [self.brightness_values[i] for i in crossings]
            plt.plot(crossings, crossing_values, 'ro', markersize=6,
                     label=f'Valid Crossings ({len(crossings)})')

        plt.xlabel('Frame Number')
        plt.ylabel('Brightness (0-1)')
        plt.title(
            f'Midline Crossing Analysis - {video_name}\nThroughput Count: {throughput_count} (Valid Crossing Pairs)')
        plt.legend()
        plt.grid(True, alpha=0.3)

        plt.tight_layout()

        # 保存图表
        os.makedirs("throughput_analysis", exist_ok=True)
        plot_path = f"throughput_analysis/{video_name}_midline_analysis.png"
        plt.savefig(plot_path, dpi=300, bbox_inches='tight')
        plt.close()

        return plot_path, throughput_count, midline

    def run_throughput_detection(self, video_path, show_video=True):
        """
        运行基于中线交点法的吞吐检测
        """
        print(f"\n🎬 开始中线交点法吞吐分析: {video_path}")
        print("📊 吞吐次数 = 有效中线交点数对")
        print("🖥️  按 'q' 键退出视频显示")

        # 重置状态
        self.brightness_values = []
        self.frame_numbers = []
        self.crossing_points = []
        self.current_throughput = 0
        self.crossing_events = []

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"❌ 无法打开视频源: {video_path}")

        # 获取视频属性
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        frame_count = 0
        processing_times = []

        # 创建进度条
        pbar = tqdm(total=total_frames, desc="处理进度", unit='帧')

        # 计算中线（初始值）
        midline = 0.5  # 初始中线值

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                frame_count += 1
                pbar.update(1)

                start_time = time.time()

                # 检测中心点
                center_point, confidence = self.detect_center_point(frame)

                brightness = 0
                if center_point:
                    # 获取中心区域亮度
                    brightness = self.get_center_region_brightness(frame, center_point)
                    self.brightness_values.append(brightness)
                    self.frame_numbers.append(frame_count)

                    # 更新中线值
                    if len(self.brightness_values) > 10:  # 至少有10个样本后再计算中线
                        midline = np.mean(self.brightness_values)

                    # 实时更新吞吐数量
                    throughput_count = self.update_throughput_count(frame_count, brightness, midline)

                processing_time = time.time() - start_time
                processing_times.append(processing_time)

                # 显示视频
                if show_video:
                    display_frame = self.draw_detection_info(frame, center_point, brightness, self.current_throughput)

                    # 调整显示窗口大小（如果视频太大）
                    display_height = min(800, height)
                    scale_factor = display_height / height
                    display_width = int(width * scale_factor)

                    resized_frame = cv2.resize(display_frame, (display_width, display_height))
                    cv2.imshow('Throughput Detection', resized_frame)

                    # 按 'q' 键退出
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        print("\n⏹️ 用户请求停止检测")
                        break

        finally:
            pbar.close()
            cap.release()
            if show_video:
                cv2.destroyAllWindows()

        # 最终的中线交点检测（更准确）
        final_throughput_count, crossings, midline = self.detect_midline_crossings()

        # 分析中线交点
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        plot_path, _, _ = self.plot_analysis(video_name)

        # 输出统计信息
        avg_processing_time = np.mean(processing_times) * 1000 if processing_times else 0

        print(f"\n📊 中线交点法检测结果:")
        print(f"总帧数: {frame_count}")
        print(f"有效亮度数据帧数: {len(self.brightness_values)}")
        print(f"中线值: {midline:.3f}")
        print(f"实时吞吐数量: {self.current_throughput}")
        print(f"最终吞吐数量: {final_throughput_count}")
        print(f"平均处理时间: {avg_processing_time:.2f}ms/帧")

        if plot_path:
            print(f"📈 分析图已保存: {plot_path}")

        return final_throughput_count

    def generate_report(self, video_path):
        """生成详细的吞吐检测报告"""
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        report_path = f"throughput_report_{video_name}.txt"

        # 检测中线交点
        throughput_count, crossings, midline = self.detect_midline_crossings()

        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("中线交点法吞吐检测报告\n")
            f.write("=" * 50 + "\n")
            f.write(f"视频名称: {video_name}\n")
            f.write(f"检测时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"总帧数: {len(self.frame_numbers)}\n")
            f.write(f"有效亮度帧数: {len(self.brightness_values)}\n")
            f.write(f"中线值: {midline:.3f}\n")
            f.write(f"检测到的有效中线交点: {len(crossings)}\n")
            f.write(f"吞吐次数 (有效交点对数): {throughput_count}\n")

            f.write("\n📊 有效中线交点事件详情:\n")
            f.write("-" * 40 + "\n")
            for i in range(0, len(crossings), 2):
                if i + 1 < len(crossings):
                    frame_num1 = self.frame_numbers[crossings[i]]
                    brightness_val1 = self.brightness_values[crossings[i]]
                    frame_num2 = self.frame_numbers[crossings[i + 1]]
                    brightness_val2 = self.brightness_values[crossings[i + 1]]

                    f.write(f"第{i // 2 + 1:03d}对交点:\n")
                    f.write(f"  交点1 - 帧号: {frame_num1:06d}, 亮度: {brightness_val1:.3f}\n")
                    f.write(f"  交点2 - 帧号: {frame_num2:06d}, 亮度: {brightness_val2:.3f}\n")
                    f.write(f"  交点间距: {abs(frame_num2 - frame_num1)} 帧\n\n")

        print(f"📝 报告已保存: {report_path}")
        return report_path


# 使用示例
if __name__ == "__main__":
    # 创建吞吐检测器
    throughput_detector = ThroughputDetector(model_path='runs/detect/train3/weights/best.pt')

    # 检测视频吞吐次数
    video_path = "D:/mkex9.mp4"

    # 运行吞吐检测（显示视频）
    throughput_count = throughput_detector.run_throughput_detection(video_path, show_video=True)

    # 生成报告
    throughput_detector.generate_report(video_path)

    print(f"\n🎯 视频 '{os.path.basename(video_path)}' 的吞吐次数: {throughput_count}")
